/**
 * Copyright (c) 2018-2023, NXOS Development Team
 * SPDX-License-Identifier: Apache-2.0
 *
 * Contains: Implementation of virtio-console
 *
 * Change Logs:
 * Date           Author            Notes
 * 2023-08-30     Shiroko           Init
 */

#include "nxos.h"

#ifdef CONFIG_NX_DRIVER_VIRTIO_CONSOLE

#undef NX_LOG_MOD_NAME
#define NX_LOG_MOD_NAME "VIRTIO-CONSOLE"

#include "base/log.h"

#define DRV_NAME "virtio console"
#define DEV_NAME "virtio_console"

#include "../virtio.h"
#include "./virtio_console.h"

NX_PRIVATE NX_Error VirtioConsoleSendCommand(NX_Device *device, NX_U32 id, NX_U16 event, NX_U16 value, void *additional)
{
    VirtioQueue *controlq = VirtioQueueGet(device, VIRTIO_CONSOLE_QUEUE_TRANSMIT_CONTROL);
    VirtioConsoleControl control = {
            .id = id,
            .event = event,
            .value = value
    };
    if (event == VIRTIO_CONSOLE_RESIZE)
    {
        if (additional == NX_NULL) return NX_EINVAL;
        NX_U32 idx[2];
        while (VirtioQueueDescAllocMany(controlq, 2, idx) != NX_EOK);
        VirtioQueueDescFill(controlq, idx[0], (NX_UArch) VirtioVirt2Phy(&control), sizeof(control),
                            VIRTIO_QUEUE_DESC_FLAG_NEXT, idx[1]);
        VirtioQueueDescFill(controlq, idx[1], (NX_UArch) VirtioVirt2Phy(additional), sizeof(NX_U16) * 2, 0, 0);
        VirtioQueuePutIntoAvailRing(controlq, idx[0]);
        VirtioSendNotification(device, VIRTIO_CONSOLE_QUEUE_TRANSMIT_CONTROL, idx[0]);
        VirtioWaitNotification(device, VIRTIO_CONSOLE_QUEUE_TRANSMIT_CONTROL, idx[0]);
        VirtioQueueDescFreeMany(controlq, idx[0]);
    }
    else
    {
        if (additional != NX_NULL) return NX_EINVAL;
        NX_U32 idx;
        while ((idx = VirtioQueueDescAlloc(controlq)) == VIRTIO_QUEUE_DESC_INVALID_ID);
        VirtioQueueDescFill(controlq, idx, (NX_UArch) VirtioVirt2Phy(&control), sizeof(control), 0, 0);
        VirtioQueuePutIntoAvailRing(controlq, idx);
        VirtioSendNotification(device, VIRTIO_CONSOLE_QUEUE_TRANSMIT_CONTROL, idx);
        VirtioWaitNotification(device, VIRTIO_CONSOLE_QUEUE_TRANSMIT_CONTROL, idx);
        VirtioQueueDescFree(controlq, idx);
    }
    return NX_EOK;
}

NX_PRIVATE NX_Error VirtioConsoleQueueNotifyHandlerForReceive(NX_Device *device, VirtioQueue *queue, NX_U32 idx);

NX_PRIVATE NX_Error VirtioConsoleOpenPort(NX_Device *device, NX_U32 portId)
{
    NX_U16 queueIndex = VIRTIO_CONSOLE_QUEUE_RECEIVE_N(portId);
    VirtioQueue *receiveq = VirtioQueueGet(device, queueIndex);
    VirtioConsoleDeviceExtension *deviceExtension = VirtioGetDevExtension(device, VirtioConsoleDeviceExtension);
    char *buffer = NX_MemAlloc(
            sizeof(char) * VIRTIO_CONSOLE_RECEIVE_BUF_SIZE * CONFIG_NX_DRIVER_VIRTIO_SPLIT_QUEUE_SIZE);
    if (buffer == NX_NULL) return NX_ENOMEM;
    NX_Fifo *queue = NX_FifoCreate(VIRTIO_CONSOLE_RECEIVE_BUF_SIZE * CONFIG_NX_DRIVER_VIRTIO_SPLIT_QUEUE_SIZE);
    if (queue == NX_NULL)
    {
        NX_MemFree(buffer);
        return NX_ENOMEM;
    }

    deviceExtension->receiveBuffers[portId] = buffer;
    deviceExtension->receiveQueues[portId] = queue;

    // Setup receive buffer
    for (NX_U32 idx = 0; idx < CONFIG_NX_DRIVER_VIRTIO_SPLIT_QUEUE_SIZE; idx++)
    {
        NX_ASSERT(receiveq->free[idx] == NX_True);
        receiveq->free[idx] = NX_False;
        receiveq->free_count -= 1;
        VirtioQueueDescFill(receiveq, idx,
                            (NX_U64) ((NX_UArch) NX_Virt2Phy(buffer + idx * VIRTIO_CONSOLE_RECEIVE_BUF_SIZE)),
                            VIRTIO_CONSOLE_RECEIVE_BUF_SIZE,
                            VIRTIO_QUEUE_DESC_FLAG_WRITE, 0);
        VirtioQueuePutIntoAvailRing(receiveq, idx);
        receiveq->trace[idx].needNotify = NX_True;
    }
    VirtioSendNotification(device, queueIndex, VIRTIO_QUEUE_DESC_INVALID_ID);
    VirtioQueueRegisterNotifier(receiveq, VirtioConsoleQueueNotifyHandlerForReceive);

    // Notify OK
    VirtioConsoleSendCommand(device, portId, VIRTIO_CONSOLE_PORT_READY, 1, NX_NULL);
    VirtioConsoleSendCommand(device, portId, VIRTIO_CONSOLE_PORT_OPEN, 1, NX_NULL);

    // select port
    deviceExtension->selectedPort = portId;

    return NX_EOK;
}

NX_PRIVATE NX_Error VirtioConsoleClosePort(NX_Device *device, NX_U32 portId)
{
    VirtioQueue *receiveq = VirtioQueueGet(device, VIRTIO_CONSOLE_QUEUE_RECEIVE_N(portId));
    VirtioQueueUnregisterNotifier(receiveq, VirtioConsoleQueueNotifyHandlerForReceive);

    VirtioConsoleDeviceExtension *deviceExtension = VirtioGetDevExtension(device, VirtioConsoleDeviceExtension);
    char *buffer = deviceExtension->receiveBuffers[portId];
    NX_Fifo *queue = deviceExtension->receiveQueues[portId];
    deviceExtension->receiveBuffers[portId] = NX_NULL;
    deviceExtension->receiveQueues[portId] = NX_NULL;
    NX_FifoDestroy(queue);
    NX_MemFree(buffer);

    return NX_EOK;
}

NX_PRIVATE NX_Error VirtioConsoleQueueNotifyHandlerForReceive(NX_Device *device, VirtioQueue *queue, NX_U32 idx)
{
    VirtioConsoleDeviceExtension *deviceExtension = VirtioGetDevExtension(device, VirtioConsoleDeviceExtension);
    NX_Size used_size = queue->used->ring[idx].len;

    NX_U32 portId = queue->queue_index < 2 ? 0 : (queue->queue_index - 2) / 2;
    char *buffer = deviceExtension->receiveBuffers[portId] + idx * VIRTIO_CONSOLE_RECEIVE_BUF_SIZE;
    NX_FifoWrite(deviceExtension->receiveQueues[portId], (void *) buffer, used_size);

    // Reuse allocated descriptor chain
    VirtioQueuePutIntoAvailRing(queue, idx);
    VirtioSendNotification(device, queue->queue_index, idx);
    return NX_EOK;
}

NX_PRIVATE NX_Error VirtioConsoleQueueNotifyHandlerForReceiveControl(NX_Device *device, VirtioQueue *queue, NX_U32 idx)
{
    VirtioConsoleDeviceExtension *deviceExtension = VirtioGetDevExtension(device, VirtioConsoleDeviceExtension);

    VirtioConsoleControl *control = (VirtioConsoleControl *) NX_Phy2Virt((NX_Addr)queue->desc[idx].addr);
    switch (control->event)
    {
    case VIRTIO_CONSOLE_DEVICE_ADD:
        NX_LOG_D("Port Added: %d.", control->id);
        deviceExtension->selectedPort = control->id;
        break;
    case VIRTIO_CONSOLE_DEVICE_REMOVE:
        NX_LOG_D("Port Removed: %d.", control->id);
        break;
    case VIRTIO_CONSOLE_CONSOLE_PORT:
        NX_LOG_D("Port nominated as console port: %d.", control->id);
        break;
    case VIRTIO_CONSOLE_PORT_OPEN:
        NX_LOG_D("Port %d status alerted: %s.", control->id, control->value == 0 ? "closed" : "opened");
        break;
    default:
        NX_LOG_W("Invalid control message received. (Id: %d, Event: 0x%x, Value: 0x%x)", control->id, control->event,
                 control->value);
        break;
    }

    // Reuse allocated descriptor chain
    VirtioQueuePutIntoAvailRing(queue, idx);
    VirtioSendNotification(device, queue->queue_index, idx);
    return NX_EOK;
}

NX_PRIVATE void VirtioConsoleCleanUp(NX_Device *device);

NX_PRIVATE NX_Error VirtioConsoleInit(NX_Driver *driver, NX_Device *device)
{
    NX_Error err = NX_EOK;
    NX_U32 features[1] = {0};
    features[0] &= VIRTIO_CONSOLE_FEATURE_SIZE | VIRTIO_CONSOLE_FEATURE_MULTIPORT;      // Enabled Features

    // Negotiate features & status
    if ((err = VirtioDeviceInitialize(device, features, 1)) != NX_EOK)
        return err;

    VirtioConsoleDeviceExtension *deviceExtension = NX_MemAlloc(sizeof(VirtioConsoleDeviceExtension));
    if (deviceExtension == NX_NULL) goto nomem_failed;
    VirtioGetExtension(device)->devExt = (void *) deviceExtension;

    VirtioConsoleDeviceConfig *deviceCfg = NX_MemAlloc(sizeof(VirtioConsoleDeviceConfig));
    if (deviceCfg == NX_NULL) goto nomem_failed;
    VirtioReadDeviceConfig(device, 0, deviceCfg, sizeof(VirtioConsoleDeviceConfig));
    deviceExtension->deviceConfig = deviceCfg;

    // Device-specified initialization
    int queue_num = 4 + ((int) deviceCfg->max_ports - 1) * 2;
    VirtioQueueAlloc(device, queue_num);
    for (int i = 0; i < queue_num; i++)
        VirtioQueueInit(device, i, CONFIG_NX_DRIVER_VIRTIO_SPLIT_QUEUE_SIZE);

    VirtioQueueRegisterNotifier(VirtioQueueGet(device, VIRTIO_CONSOLE_QUEUE_RECEIVE_CONTROL),
                                VirtioConsoleQueueNotifyHandlerForReceiveControl);

    deviceExtension->receiveQueues = NX_MemAlloc(sizeof(NX_Fifo *) * deviceCfg->max_ports);
    if (deviceExtension->receiveQueues == NX_NULL) goto nomem_failed;
    NX_MemSet(deviceExtension->receiveQueues, 0, sizeof(NX_Fifo *) * deviceCfg->max_ports);
    deviceExtension->receiveBuffers = NX_MemAlloc(sizeof(char *) * deviceCfg->max_ports);
    if (deviceExtension->receiveBuffers == NX_NULL) goto nomem_failed;
    NX_MemSet(deviceExtension->receiveBuffers, 0, sizeof(NX_Fifo *) * deviceCfg->max_ports);

    // Driver OK
    VirtioDeviceDriverOK(device);

    deviceExtension->selectedPort = 0;
    deviceExtension->receiveControlBuffer = NX_MemAlloc(
            sizeof(VirtioConsoleControl) * CONFIG_NX_DRIVER_VIRTIO_SPLIT_QUEUE_SIZE);
    if (deviceExtension->receiveControlBuffer == NX_NULL) goto nomem_failed;

    // Populate controlq of rx
    VirtioQueue *controlq = VirtioQueueGet(device, VIRTIO_CONSOLE_QUEUE_RECEIVE_CONTROL);
    for (NX_U32 idx = 0; idx < CONFIG_NX_DRIVER_VIRTIO_SPLIT_QUEUE_SIZE; idx++)
    {
        NX_ASSERT(controlq->free[idx] == NX_True);
        controlq->free[idx] = NX_False;
        controlq->free_count -= 1;
        VirtioQueueDescFill(controlq, idx,
                            (NX_U64) ((NX_UArch) NX_Virt2Phy(&deviceExtension->receiveControlBuffer[idx])),
                            sizeof(VirtioConsoleControl),
                            VIRTIO_QUEUE_DESC_FLAG_WRITE, 0);
        VirtioQueuePutIntoAvailRing(controlq, idx);
        controlq->trace[idx].needNotify = NX_True;
    }

    VirtioSendNotification(device, VIRTIO_CONSOLE_QUEUE_RECEIVE_CONTROL, VIRTIO_QUEUE_DESC_INVALID_ID);
    VirtioConsoleSendCommand(device, 0, VIRTIO_CONSOLE_DEVICE_READY, 1, NX_NULL);

    return NX_EOK;

nomem_failed:
    err = NX_ENOMEM;

    VirtioConsoleCleanUp(device);
    return err;
}

NX_PRIVATE void VirtioConsoleCleanUp(NX_Device *device)
{
    // Reset the device
    VirtioDeviceReset(device);
    // Cleanup VirtioQueue
    for (int i = 0; i < VirtioGetExtension(device)->queue_num; i++)
        VirtioQueueDestroy(device, i);
    VirtioQueueFree(device);

    VirtioConsoleDeviceExtension *deviceExtension = VirtioGetDevExtension(device, VirtioConsoleDeviceExtension);
    if (deviceExtension == NX_NULL) return;
    if (deviceExtension->deviceConfig)
    {
        NX_U32 ports = deviceExtension->deviceConfig->max_ports;
        NX_MemFree(deviceExtension->deviceConfig);

        if (deviceExtension->receiveBuffers)
        {
            for (int i = 0; i < ports; i++)
                if (deviceExtension->receiveBuffers[i] != NX_NULL)
                    NX_MemFree(deviceExtension->receiveBuffers[i]);
            NX_MemFree(deviceExtension->receiveBuffers);
        }

        if (deviceExtension->receiveQueues)
        {
            for (int i = 0; i < ports; i++)
                if (deviceExtension->receiveQueues[i] != NX_NULL)
                    NX_MemFree(deviceExtension->receiveQueues[i]);
            NX_MemFree(deviceExtension->receiveQueues);
        }

        if (deviceExtension->receiveControlBuffer)
            NX_MemFree(deviceExtension->receiveControlBuffer);
    }

    NX_MemFree(deviceExtension);
}

NX_PRIVATE NX_Error open(struct NX_Device *device, NX_U32 flags)
{
    VirtioConsoleDeviceExtension *deviceExtension = VirtioGetDevExtension(device, VirtioConsoleDeviceExtension);
    VirtioConsoleOpenPort(device, deviceExtension->selectedPort);
    return NX_EOK;
}

NX_PRIVATE NX_Error close(struct NX_Device *device)
{
    VirtioConsoleDeviceExtension *deviceExtension = VirtioGetDevExtension(device, VirtioConsoleDeviceExtension);
    VirtioConsoleClosePort(device, deviceExtension->selectedPort);
    return NX_EOK;
}

NX_PRIVATE NX_Error read(struct NX_Device *device, void *buf, NX_Offset off, NX_Size len, NX_Size *outLen)
{
    VirtioConsoleDeviceExtension *deviceExtension = VirtioGetDevExtension(device, VirtioConsoleDeviceExtension);
    NX_Fifo *fifo = NX_NULL;
    if ((fifo = deviceExtension->receiveQueues[deviceExtension->selectedPort]) != NX_NULL)
    {
        *outLen = NX_FifoRead(fifo, buf, len);
        return NX_EOK;
    }
    return NX_EINVAL;
}

NX_PRIVATE NX_Error write(struct NX_Device *device, void *buf, NX_Offset off, NX_Size len, NX_Size *outLen)
{
    VirtioConsoleDeviceExtension *deviceExtension = VirtioGetDevExtension(device, VirtioConsoleDeviceExtension);

    NX_U16 queueIndex = VIRTIO_CONSOLE_QUEUE_TRANSMIT_N(deviceExtension->selectedPort);
    VirtioQueue *transmitq = VirtioQueueGet(device, queueIndex);
    NX_U32 idx;
    while ((idx = VirtioQueueDescAlloc(transmitq)) == VIRTIO_QUEUE_DESC_INVALID_ID);
    VirtioQueueDescFill(transmitq, idx, (NX_UArch) VirtioVirt2Phy(buf), len, 0, 0);
    VirtioQueuePutIntoAvailRing(transmitq, idx);
    VirtioSendNotification(device, queueIndex, idx);
    VirtioWaitNotification(device, queueIndex, idx);
    VirtioQueueDescFree(transmitq, idx);

    return NX_EINVAL;
}

NX_PRIVATE NX_Error control(struct NX_Device *device, NX_U32 cmd, void *arg)
{
    return NX_EOK;
}

NX_PRIVATE NX_Error poll(struct NX_Device *device, NX_PollState *pState)
{
    VirtioConsoleDeviceExtension *deviceExtension = VirtioGetDevExtension(device, VirtioConsoleDeviceExtension);
    NX_Fifo *fifo = NX_NULL;
    if ((fifo = deviceExtension->receiveQueues[deviceExtension->selectedPort]) != NX_NULL)
        return NX_FifoPoll(fifo, pState);
    return NX_EINVAL;
}


NX_PRIVATE NX_DriverOps virtioConsoleDriverOps = {
        .open = open,
        .close = close,
        .read = read,
        .write = write,
        .control = control,
        .poll = poll
};


NX_PRIVATE VirtioDriverTableEntity virtioConsoleDriverTableEntity = {
        .deviceId = VIRTIO_DEVICE_ID_CONSOLE,
        .driverName = "virtio-console",
        .deviceNameTemplate = "virtio_console_%d",
        .deviceType = NX_DEVICE_TYPE_UART,
        .driverOps = &virtioConsoleDriverOps,
        .initHandler = VirtioConsoleInit,
        .cleanUpHandler = VirtioConsoleCleanUp
};

VirtioDriverTableEntity *VirtioDriverTableEntityForVirtioConsole(void)
{
    return &virtioConsoleDriverTableEntity;
}

#else

#include "../virtio.h"

VirtioDriverTableEntity *VirtioDriverTableEntityForVirtioConsole(void)
{
    return NX_NULL;
}

#endif // CONFIG_NX_DRIVER_VIRTIO_CONSOLE
